#include "LeNetINT8.h"
#include "ParamsLoader.h"
#include "BasicOp.h"
#include "QAdvanceOp.h"
#include "QBasicOp.h"
#include "Statistics.h"

using namespace DL;

const std::string lables[] = { "plane", "car", "bird", "cat", "deer","dog", "frog", "horse", "ship", "truck" };

LeNetINT8::LeNetINT8(std::string paramsFile)
{
	ParamsLoader paramsLoader;
	paramsLoader.LoadFromFile(paramsFile);

	this->conv1Weight = &paramsLoader.LoadTensor4DInt8("conv1.weightQTA");
	this->conv1WeightScale = &paramsLoader.LoadTensor1DFloat("conv1.scale");
	this->conv1Bias = &paramsLoader.LoadTensor1DFloat("conv1.bias");
	this->conv1OutputScale = paramsLoader.LoadfloatValue("conv1.output.scale");
	this->conv1OutputZero = paramsLoader.LoadIntValue("conv1.output.zero");

	this->conv2Weight = &paramsLoader.LoadTensor4DInt8("conv2.weightQTA");
	this->conv2WeightScale = &paramsLoader.LoadTensor1DFloat("conv2.scale");
	this->conv2Bias = &paramsLoader.LoadTensor1DFloat("conv2.bias");
	this->conv2OutputScale = paramsLoader.LoadfloatValue("conv2.output.scale");
	this->conv1OutputZero = paramsLoader.LoadIntValue("conv2.output.zero");

	this->linear1Weight = &paramsLoader.LoadTensor2DInt8("fc1.weightQTA");
	this->linear1WeightScale = &paramsLoader.LoadTensor1DFloat("fc1.scale");
	this->linear1Bias = &paramsLoader.LoadTensor1DFloat("fc1.bias");
	this->linear1OutputScale = paramsLoader.LoadfloatValue("fc1.output.scale");
	this->linear1OutputZero = paramsLoader.LoadIntValue("fc1.output.zero");

	this->linear2Weight = &paramsLoader.LoadTensor2DInt8("fc2.weightQTA");
	this->linear2WeightScale = &paramsLoader.LoadTensor1DFloat("fc2.scale");
	this->linear2Bias = &paramsLoader.LoadTensor1DFloat("fc2.bias");
	this->linear2OutputScale = paramsLoader.LoadfloatValue("fc2.output.scale");
	this->linear2OutputZero = paramsLoader.LoadIntValue("fc3.output.zero");

	this->linear3Weight = &paramsLoader.LoadTensor2DInt8("fc3.weightQTA");
	this->linear3WeightScale = &paramsLoader.LoadTensor1DFloat("fc3.scale");
	this->linear3Bias = &paramsLoader.LoadTensor1DFloat("fc3.bias");
	this->linear3OutputScale = paramsLoader.LoadfloatValue("fc3.output.scale");
	this->linear3OutputZero = paramsLoader.LoadIntValue("fc3.output.zero");
}

LeNetINT8::~LeNetINT8()
{
	delete this->conv1Weight;
	delete this->conv1WeightScale;
	delete this->conv1Bias;
	delete this->conv2Weight;
	delete this->conv2WeightScale;
	delete this->conv2Bias;
	delete this->linear1Weight;
	delete this->linear1WeightScale;
	delete this->linear1Bias;
	delete this->linear2Weight;
	delete this->linear2WeightScale;
	delete this->linear2Bias;
	delete this->linear3Weight;
	delete this->linear3WeightScale;
	delete this->linear3Bias;
}

int LeNetINT8::GetResultIndex(Image& image)
{
	DL::Tensor3D<float> imageTensor = DL::Tensor3D<float>(image);
	float mean[] = { 0.5f,0.5f,0.5f };
	float std[] = { 0.5f,0.5f,0.5f };
	Tensor3D<float>& input = Normalize(imageTensor, mean, std);
	Tensor3D<int8>& inputQat = Quantify(input, { 1.0f / 127 ,0 });
	Tensor3D<int8>& conv1Result = QConv(inputQat, { 1.0f / 127 ,0 }, 1, 0, 2, *conv1Weight, *conv1WeightScale, *conv1Bias, { conv1OutputScale ,0 });
	Tensor3D<int8>& conv2Result = QConv(conv1Result, { conv1OutputScale ,0 }, 1, 0, 2, *conv2Weight, *conv2WeightScale, *conv2Bias, { conv2OutputScale ,0 });
	Tensor1D<int8>& flattenResult = QFlatten(conv2Result);
	Tensor1D<int8>& linear1Result = QFullConnect(flattenResult, { conv2OutputScale ,0 }, *linear1Weight, *linear1WeightScale, *linear1Bias, { linear1OutputScale ,0 });
	Tensor1D<int8>& linear2Result = QFullConnect(linear1Result, { linear1OutputScale ,0 }, *linear2Weight, *linear2WeightScale, *linear2Bias, { linear2OutputScale ,0 });
	Tensor1D<int8>& linear3Result = QLinear(linear2Result, { linear2OutputScale ,0 }, *linear3Weight, *linear3WeightScale, *linear3Bias, { linear3OutputScale ,0 });

	//int8 minv, maxv;
	//GetRange(linear3Result, maxv, minv);
	//std::cout << (int)minv << ":" << (int)maxv << std::endl;

	Tensor1D<float>& resultDeqta = Dequantify(linear3Result, { linear3OutputScale ,0 });
	Tensor1D<float>& result = SoftMax(resultDeqta);

	float max = -FLT_MAX;
	int index = -1;
	for (int i = 0; i < result.width; i++)
	{
		float value = result.Get(i);
		if (value > max) {
			max = value;
			index = i;
		}
	}
	delete& input;
	delete& inputQat;
	delete& conv1Result;
	delete& conv2Result;
	delete& flattenResult;
	delete& linear1Result;
	delete& linear2Result;
	delete& linear3Result;
	delete& resultDeqta;
	delete& result;
	return index;
}

std::string LeNetINT8::GetResultLable(Image& image)
{
	return lables[GetResultIndex(image)];
}
